package com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.filters;

import com.apache.tools.StrUtil;
import com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.config.Rate;
import com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.config.RateLimitKeyGenerator;
import com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.config.RateLimitUtils;
import com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.config.RateLimiter;
import com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.config.properties.RateLimitProperties;

import java.util.concurrent.TimeUnit;

import com.marcosbarbero.cloud.autoconfigure.zuul.ratelimit.support.RateLimitExceededException;
import com.netflix.zuul.context.RequestContext;
import org.springframework.cloud.netflix.zuul.filters.Route;
import org.springframework.cloud.netflix.zuul.filters.RouteLocator;
import org.springframework.http.HttpStatus;
import org.springframework.web.util.UrlPathHelper;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

/**
 * @description:
 * @author: Hou Dayu
 * @date: Created in 2020/9/12
 */
public class RateLimitPreFilter extends AbstractRateLimitFilter {

    private final RateLimitProperties properties;
    private final RateLimiter rateLimiter;
    private final RateLimitKeyGenerator rateLimitKeyGenerator;

    public RateLimitPreFilter(RateLimitProperties properties, RouteLocator routeLocator, UrlPathHelper urlPathHelper, RateLimiter rateLimiter, RateLimitKeyGenerator rateLimitKeyGenerator, RateLimitUtils rateLimitUtils) {
        super(properties, routeLocator, urlPathHelper, rateLimitUtils);
        this.properties = properties;
        this.rateLimiter = rateLimiter;
        this.rateLimitKeyGenerator = rateLimitKeyGenerator;
    }

    public String filterType() {
        return "pre";
    }

    public int filterOrder() {
        return this.properties.getPreFilterOrder();
    }

    public Object run() {
        RequestContext ctx = RequestContext.getCurrentContext();
        HttpServletResponse response = ctx.getResponse();
        HttpServletRequest request = ctx.getRequest();
        Route route = this.route(request);
        this.policy(route, request).forEach((policy) -> {
            String key = this.rateLimitKeyGenerator.key(request, route, policy);
            if(StrUtil.isNotNull(key)) {
                Rate rate = this.rateLimiter.consume(policy, key, (Long) null);
                String httpHeaderKey = key.replaceAll("[^A-Za-z0-9-.]", "_").replaceAll("__", "_");
                Long limit = policy.getLimit();
                Long remaining = rate.getRemaining();
                if (limit != null) {
                    response.setHeader("X-RateLimit-Limit-" + httpHeaderKey, String.valueOf(limit));
                    response.setHeader("X-RateLimit-Remaining-" + httpHeaderKey, String.valueOf(Math.max(remaining, 0L)));
                }
                Long quota = policy.getQuota();
                Long remainingQuota = rate.getRemainingQuota();
                if (quota != null) {
                    request.setAttribute("rateLimitRequestStartTime", System.currentTimeMillis());
                    response.setHeader("X-RateLimit-Quota-" + httpHeaderKey, String.valueOf(quota));
                    response.setHeader("X-RateLimit-Remaining-Quota-" + httpHeaderKey, String.valueOf(TimeUnit.MILLISECONDS.toSeconds(Math.max(remainingQuota, 0L))));
                }
                response.setHeader("X-RateLimit-Reset-" + httpHeaderKey, String.valueOf(rate.getReset()));
                if (limit != null && remaining < 0L || quota != null && remainingQuota < 0L) {
                    ctx.setResponseStatusCode(HttpStatus.TOO_MANY_REQUESTS.value());
                    ctx.put("rateLimitExceeded", "true");
                    ctx.setSendZuulResponse(false);
                    throw new RateLimitExceededException();
                }
            }
        });
        return null;
    }
}
